import asyncio
import json
from contextlib import AsyncExitStack
from typing import Optional
from mcp import StdioServerParameters, ClientSession
from mcp.client.stdio import stdio_client
from openai import OpenAI
from openai.types.chat.chat_completion import Choice

from connect_factory import *
class MCPClient:

    def __init__(self,manager:ServerManager):
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )
        self.server_manager = manager

    async def execute(self,question:str):
        # 一、通过session列表所有的工具 及组装function_calling
        functions_dict = self.server_manager.mcp_functions
        tools = []
        for key,mcp_function in functions_dict.items():
            #转换成function calling 格式
            tools.append(
            {
             "type":"function",
             "function":{
                 "name":mcp_function.name,
                 "description":mcp_function.description,
                 "parameters":mcp_function.input_schema
                }
             } )


        # 二、 大模型调用参数：组装messages
        messages = [
            {
                "role": "user",
                "content": question
            }
        ]
        # 调用大模型
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=tools,
        )
        # 获取大模型的决策结果
        print("==== deepseek 决策结果：",deepseek_response)
        choice_result = deepseek_response.choices[0]
        print("最终选择的工具：",choice_result.message.tool_calls[0].function.name)
        #迭代调用，直到stop为止
        result = await self.large_model_strategy(choice_result,messages,tools,1)
        print("*** Large Model最终结果：",result)

    #大模型做决策——工具链路执行
    async def large_model_strategy(self,choice_result:Choice,messages:list[dict],tools:list[dict],cnt)->str:
        print(" ==== 调用迭代函数large_model_strategy ===>第",cnt,"次调用 ====")
        cnt += 1
        '''  ******************* 3.4 function calling 过程 *******************'''
        if choice_result.finish_reason == "stop":
            #print("最终结果：",choice_result.message.content)
            return choice_result.message.content

        if choice_result.finish_reason == "tool_calls":
            # 根据大模型选择的工具，进行调用执行
            messages.append(choice_result.message.model_dump())
            # 调用工具链
            tool_calls = choice_result.message.tool_calls
            for tool_call in tool_calls:
                function_name = tool_call.function.name
                arguments = json.loads(tool_call.function.arguments)
                # 调用工具链
                call_result = await self.server_manager.execute_calling(function_name, arguments)
                print(f"==== {function_name} 执行结果：", call_result)
                # 保存到消息列表中，大模型使用的结果参数
                messages.append({
                    "role": "tool",
                    "content": call_result,
                    "tool_call_id": tool_call.id
                })

            ''' ************** 推理结果 (多次调用大模型）*************'''

            # 再次调用大模型
            deepseek_response = self.deepseek.chat.completions.create(
                model="deepseek-chat",
                messages=messages,
                tools=tools,
            )
            choice_result = deepseek_response.choices[0]
            #迭代调用，直到stop为止
            return await self.large_model_strategy(choice_result,messages,tools,cnt)

async def main():
    manager = ServerManager()
    client = MCPClient(manager)
    try:
        #初始化资源
        await manager.initialize()
        #从键盘输入问题
        while True:
            question = input("请输入你的问题：")
            if question.lower() == "exit":
                break
            await client.execute(question)

    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await manager.aclose()
        #延迟关闭
        await asyncio.sleep(0.1)

if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    try:
        loop.run_until_complete(main())
    finally:
        loop.close()

